# coding: utf-8

from typing import IO, Sequence, Optional
from graph import UVEdge, vertex_t, Subgraph, Hierarchy
from utils import sgn
from io import StringIO
import copy
import time
import math
from .Layering import Layering


class TreeEdge:
    u: vertex_t
    v: vertex_t
    dir: int

    def __init__(self, u: vertex_t, v: vertex_t, dir: int):
        self.u = u
        self.v = v
        self.dir = dir

    def __str__(self):
        return f"{UVEdge(self.u, self.v) | self.dir}"


class TreeNode:
    parent: Optional[vertex_t] = None
    u: vertex_t = -1

    cut_value: int = 0
    out_cut_value: int = 0

    children: Sequence[vertex_t]

    # for marking the order of nodes using pre-order traversal
    min: int = 10000
    order: int = -1

    def __init__(self):
        self.children = []  # 必须在__init__里初始化, 否则会共用同一个list

    # def __deepcopy__(self, memodict={}):
    #     obj = TreeNode()
    #     obj.parent = self.parent
    #     obj.u = self.u
    #     obj.cut_value = self.cut_value
    #     obj.out_cut_value = self.out_cut_value
    #     obj.children = self.children[:]
    #     obj.min = self.min
    #     obj.order = self.order
    #     return obj

    def __str__(self):
        out = StringIO()
        out.write(
            f'u:{self.u} parent:{self.parent} children:[{" ".join([str(v) for v in self.children])}] order:{self.order}\n'
        )
        return out.getvalue()


class TightTree:
    root: vertex_t
    h: Hierarchy

    def __init__(self, h: Hierarchy, root: vertex_t):
        self.root = root
        self.h = h
        g = h.graph()
        # g.set_all_nodes_attr('nodes_tree', TreeNode())
        for u in g.get_all_nodes():
            u_node = TreeNode()
            u_node.u = u
            # print(u, id(u_node), id(u_node.children))
            g.set_node_attr(u, "nodes_tree", u_node)

    def get_node(self, u: vertex_t) -> TreeNode:
        g = self.h.graph()
        return g.get_node_attr(u, "nodes_tree")

    def get_cut_val(self, u: vertex_t, v: vertex_t) -> int:
        assert u == self.get_node(v).parent
        return self.get_node(v).cut_value

    def set_cut_val(self, u: vertex_t, v: vertex_t, val: int):
        assert u == self.get_node(v).parent
        self.get_node(v).cut_value = val

    def get_out_cut_val(self, u: vertex_t, v: vertex_t) -> int:
        assert u == self.get_node(v).parent
        return self.get_node(v).out_cut_value

    def set_out_cut_val(self, u: vertex_t, v: vertex_t, val: int):
        assert u == self.get_node(v).parent
        self.get_node(v).out_cut_value = val

    # AM I SURE ABOUT THIS?

    def add_child(self, parent: vertex_t, child: vertex_t):
        self.get_node(parent).children.append(child)
        self.get_node(child).parent = parent

    def remove_child(self, parent: vertex_t, child: vertex_t):
        self.unlink_child(parent, child)
        self.get_node(child).parent = -1

    def unlink_child(self, parent: vertex_t, child: vertex_t):
        try:
            children = self.get_node(parent).children
            i = children.index(child)
            del children[i]
        except Exception:
            pass

    # END SURE

    """
     Find which component of the tree <u> falls in, if the tree is split by removing <e>.
     Postorder **must*by calculated first.
    """

    def component(self, e: UVEdge, u: vertex_t) -> vertex_t:
        to_node = self.get_node(e.v)
        u_node = self.get_node(u)
        if to_node.min > u_node.min or to_node.order < u_node.order:
            return e.u
        return e.v

    # 1 if (parent, u) is an edge, -1 if (u, parent) is an edge

    def dir(self, u: vertex_t, v: vertex_t) -> int:
        return sgn(self.h.span(u, v))

    # find a common ancestor of two vertices, postorder **must*be calculated first
    def common_acestor(self, u: vertex_t, v: vertex_t) -> vertex_t:
        ancestor: vertex_t = u

        left_min: int = -1
        right_order: int = -1
        u_node = self.get_node(u)
        v_node = self.get_node(v)

        if u_node.order < v_node.order:
            left_min = u_node.min
            right_order = v_node.order
        else:
            left_min = v_node.min
            right_order = u_node.order
        # the common acestor is the first node such that min <= left_min  order]= right_order
        ancestor_node = self.get_node(ancestor)
        while ancestor_node.min > left_min or ancestor_node.order < right_order:
            # no need to check if parent exists, root has the smalles min and the largest order
            ancestor = ancestor_node.parent
            ancestor_node = self.get_node(ancestor)
        return ancestor

    # template<typename F>
    # void for_each(F f, vertex_t u) const {
    #     f(u);
    #     for ( const auto& child : children(u) ) {
    #         for_each(f, child);
    #     }
    # }

    """
     Performs a postorder search of the subtree rooted at <u>.
     The first leaf node is given the order <order>.
    """

    def postorder_search(self, u: vertex_t, order: int) -> int:
        u_node = self.get_node(u)
        u_node.min = order
        for child in u_node.children:
            order = self.postorder_search(child, order)

        u_node.order = order
        return order + 1

    def _print(self, u: vertex_t, depth: int, out: IO[str]) -> str:
        indent: int = 4
        for _ in range(depth):
            for _ in range(indent):
                out.write(" ")

        out.write(f"{u}(")
        u_node = self.get_node(u)
        out.write(f"dir={ 0 if u == self.root else self.dir(u_node.parent, u) } ")
        out.write(f"cut={u_node.cut_value} ")
        out.write(f"out={u_node.out_cut_value} ")
        out.write(f"ord={u_node.order} ")
        out.write(f"min={u_node.min} ")
        out.write(")\n")

        for child in u_node.children:
            self._print(child, depth + 1, out)

    def __str__(self):
        out = StringIO()
        self._print(self.root, 0, out)
        return out.getvalue()

    """
     Swaps a non-tree edge <entering> for a tree edge <leaving> and reclassures the tree so
     that <leaving.to> is the predecessor of all edges in the subtree rooted at <leaving.to>
    """
    # TODO

    def swap_edges(self, entering: UVEdge, leaving: UVEdge):
        parent: vertex_t = entering.u
        u: vertex_t = entering.v
        while u != leaving.u:
            # std::cout << u << "\n"   print(u,v)
            u_node = self.get_node(u)
            p_node = self.get_node(parent)
            tmp: vertex_t = u_node.parent
            u_node.parent = parent
            p_node.children.append(u)
            parent = u
            u = tmp
            self.unlink_child(u, parent)


"""
 Network simple algorithm for Layering a graph.
"""


class NetworkSimpleLayering(Layering):
    tree: TightTree

    def run(self, g: Subgraph):
        if g.node_count() == 0:
            return Hierarchy(g)

        h = self.init_Hierarchy(g)

        self.init_tree(h)
        self.init_cut_values()
        self.optimize_edge_length(g, h)

        # Hierarchy h = initialize_ranking(g)
        # initialize_tree(g, h)
        # postorder_search()
        # init_cut_values(g, h)

        # #std::cout << h << "\n\n"
        # #std::cout << tree << "\n"

        # optimize_edge_length(g, h)

        min: int = math.inf
        max: int = -1 * math.inf
        for u in g.get_all_nodes():
            if h.get_rank(u) > max:
                max = h.get_rank(u)
            if h.get_rank(u) < min:
                min = h.get_rank(u)

        h.layers = [[] for _ in range(max - min + 1)]

        for u in g.get_all_nodes():
            h.set_rank(u, h.get_rank(u) - min)
            h.layers[h.get_rank(u)].append(u)
            h.set_pos(u, len(h.layers[h.get_rank(u)]) - 1)

        # int total = 0
        # for (auto u : h.g.get_all_nodes()):
        #     for (auto v : h.g.out_neighbours(u)):
        #         total += h.span(u, v)

        # std::cout << total << "\n"

        # std::cout << "done\n"

        # print("----------------------未进行节点排序之前的分层", h.layers)
        total_elements = sum(len(sublist) for sublist in h.layers)
        # print(total_elements)
        # 计算时间消耗，特别是对程序运行时间消耗
        start = time.time()

        copy_h_layers = copy.deepcopy(h.layers)
        new_layers = []
        new_layers.append(copy_h_layers[0])
        for i in range(len(copy_h_layers) - 1):
            c = []
            for j in range(len(copy_h_layers[i])):

                c.append(h.g.out_neighbours(copy_h_layers[i][j]))
            d = []
            for k in range(len(c)):
                for j in range(len(c[k])):
                    d.append(c[k][j])
            # print('------------------------', d)

            seen = {}
            result = []
            for idx, num in reversed(list(enumerate(d))):
                if num not in seen:
                    seen[num] = idx
                    result.append(num)
            e = list(reversed(result))
            copy_h_layers[i + 1] = e

            new_layers.append(copy_h_layers[i + 1])

        # print(new_layers)
        end = time.time()
        # print("循环运行时间:%.10f秒" % (end - start))

        return h

    """
     Assignes each vertex a layer, such that each edge goes from a lower layer to higher one
     and the source vertices are at the lowest layer.
     The resulting rank is not normalized, meaning that the lowest layer doesnt have to be 0.
     It is also not neccesarly the optimal rank.
     @param g the graph whose vertices are to be assigned to layers
     @return resulting Hierarchy, only the rank of nodes is defined, layers and pos are undefined
    """

    def init_Hierarchy(self, g: Subgraph) -> Hierarchy:
        h = Hierarchy(g, -1)
        curr: int = 0
        processed: int = 0

        to_rank: Sequence[vertex_t] = []
        while processed < g.node_count():
            to_rank = []
            for u in g.get_all_nodes():
                if h.get_rank(u) == -1:
                    # check if there are any edges going from unranked vertices
                    source: bool = True
                    for v in g.in_neighbours(u):
                        if h.get_rank(v) == -1:
                            source = False
                            break

                    if source:
                        to_rank.append(u)
                        processed += 1

            for u in to_rank:
                h.set_rank(u, curr)
            curr += 1

        return h

    """
     Conclasss a tree of all vertices
     reachable from the root through tight edges.
    """

    def basic_tree(self, h: Hierarchy, root: vertex_t) -> int:
        g = h.graph()
        g.set_node_attr(root, "done", True)
        added: int = 1
        # print('basic_tree:', root, g.neighbours(root))
        for u in g.neighbours(root):
            span = h.span(root, u)
            if (not g.get_node_attr(u, "done")) and abs(span) == 1:
                self.tree.add_child(root, u)
                # print('add child:', u)
                added += self.basic_tree(h, u)

        return added

    """
     Finds a spanning tree of tight edges.
     Tight edge is any edge (u, v) for which rank[v] - rank[u] == 1.
     @param g the graph
     @param rank rank of the vertices of the graph
     @return the spanning tree
    """

    def init_tree(self, h: Hierarchy):
        g = h.graph()
        # TODO get_node(0)是否第一个节点?
        self.tree = TightTree(h, g.get_node(0))
        g.set_all_nodes_attr("done", False)

        finished = self.basic_tree(h, self.tree.root)
        # for v in g.g.g.vs:
        #     print(v['nodes'])
        # print('init_tree:', self.tree)
        # print('====================')
        while finished < g.node_count():
            # in the underlying undirected graph find edge (u, v) with the smallest span
            # such that u is already in the tree and v is not in the tree
            e = UVEdge(-1, -1)
            min_span = math.inf
            for u in g.get_all_nodes():
                if g.get_node_attr(u, "done"):
                    for v in g.neighbours(u):
                        span = h.span(u, v)
                        if (not g.get_node_attr(v, "done")) and abs(span) < abs(
                            min_span
                        ):
                            e = UVEdge(u, v)
                            min_span = span
            # print('=======================--=-=-=-=', finished, g.node_count())

            # make the edge tight by moving all the vertices in the tree
            dir = sgn(min_span)
            min_span -= dir
            for u in g.get_all_nodes():
                if g.get_node_attr(u, "done"):
                    h.set_rank(u, h.get_rank(u) + min_span)

            self.tree.add_child(e.u, e.v)
            g.set_node_attr(e.v, "done", True)
            finished += 1
        self.tree.postorder_search(self.tree.root, 0)

        # print('after init tree:', self.tree)

    # Calculates the initial cut values of all edges in the tight tree.

    def init_cut_values(self):
        root_node = self.tree.get_node(self.tree.root)
        for child in root_node.children:
            self._init_cut_values(self.tree.root, child)

    # Calculates cut values of all edges in the subtree rooted at <v> and then the cut value of (<u>, <v>).

    def _init_cut_values(self, u: vertex_t, v: vertex_t):
        v_node = self.tree.get_node(v)
        for child in v_node.children:
            self._init_cut_values(v, child)

        self.set_cut_value(u, v)

    """
     Calculates cut value of the UVEdge (u, v).
     Requires the cut values of the edges between <v> and its children to be already calculated.
    """

    def set_cut_value(self, u: vertex_t, v: vertex_t):
        # print("set cut value: ", UVEdge(u, v), "| ", self.tree.node(v).dir)
        g = self.tree.h.graph()
        val: int = 0

        v_node = self.tree.get_node(v)
        for child in v_node.children:
            val += (
                self.tree.dir(u, v)
                * self.tree.dir(v, child)
                * self.tree.get_out_cut_val(v, child)
            )
            # print("after child ", val)

        for x in g.neighbours(v):
            # print("fixing ", v, ": ", x)
            if self.tree.component(UVEdge(u, v), x) == u:
                val += self.tree.dir(x, v) * self.tree.dir(u, v)
                # print("after ", UVEdge(x, v), ": ", val)

        self.tree.set_cut_val(u, v, val)

        for x in g.neighbours(u):
            # print("fixing ", u, ": ",  x)
            if self.tree.component(UVEdge(u, v), x) == v:
                val -= self.tree.dir(u, x) * self.tree.dir(u, v)

        self.tree.set_out_cut_val(u, v, val)

    def switch_tree_edges(
        self, g: Subgraph, h: Hierarchy, orig: TreeEdge, entering: TreeEdge
    ):
        ancestor = self.tree.common_acestor(entering.u, entering.v)

        self.tree.swap_edges(UVEdge(entering.u, entering.v), UVEdge(orig.u, orig.v))

        self.tree.postorder_search(ancestor, self.tree.get_node(ancestor).min)
        self.fix_cut_values(g, h, ancestor, orig.u)
        self.fix_cut_values(g, h, ancestor, orig.v)

    def fix_cut_values(self, g: Subgraph, h: Hierarchy, root: vertex_t, u: vertex_t):
        # print("fix cut values from ", u, " with root ", root)
        while u != root:
            # print("u: ", u)
            parent: vertex_t = self.tree.get_node(u).parent
            self.set_cut_value(parent, u)
            u = parent

    def move_subtree(self, h: Hierarchy, root: vertex_t, d: int):
        g = h.graph()
        g.set_node_attr(root, "rank", g.get_node_attr(root, "rank") + d)
        root_node = self.tree.get_node(root)
        for child in root_node.children:
            self.move_subtree(h, child, d)

    def find_entering_edge(
        self, g: Subgraph, h: Hierarchy, leaving: TreeEdge
    ) -> TreeEdge:
        entering = TreeEdge(0, 0, -leaving.dir)
        span = math.inf

        start_component: vertex_t = leaving.v if leaving.dir == 1 else leaving.u
        end_component: vertex_t = (
            leaving.v if start_component == leaving.u else leaving.u
        )

        for u in g.get_all_nodes():
            if self.tree.component(UVEdge(leaving.u, leaving.v), u) != start_component:
                continue

            for v in g.out_neighbours(u):
                if (
                    self.tree.component(UVEdge(leaving.u, leaving.v), v)
                    == end_component
                    and h.span(u, v) < span
                ):
                    entering.u = u if entering.dir == 1 else v
                    entering.v = v if entering.dir == 1 else u
                    span = h.span(u, v)

        assert span < math.inf
        return entering

    def find_leaving_edge(self, h: Hierarchy) -> Optional[TreeEdge]:
        g = h.graph()
        for u in g.get_all_nodes():
            for v in g.out_neighbours(u):
                v_node = self.tree.get_node(v)
                if v != self.tree.root and v_node.parent == u and v_node.cut_value < 0:
                    # print('--', v, self.tree.root,
                    #       v_node.parent, v_node.cut_value)
                    return TreeEdge(u, v, 1)

                u_node = self.tree.get_node(u)
                if u != self.tree.root and u_node.parent == v and u_node.cut_value < 0:
                    # print('==', u, self.tree.root,
                    #       v_node.parent, v_node.cut_value)
                    return TreeEdge(v, u, -1)

        return None

    # Applies a function to each node of the subtree with given root in a preorder search
    # def for_each_node(self, root: vertex_t, f: Callable[[vertex_t], Any]):
    #     f(root)
    #     root_node = self.tree.get_node(root)
    #     for child in root_node.children:
    #         self.for_each_node(child, f)

    def optimize_edge_length(self, g: Subgraph, h: Hierarchy):
        iters: int = 0

        while True:
            leaving = self.find_leaving_edge(h)
            if leaving is None:
                break

            entering = self.find_entering_edge(g, h, leaving)
            # print('-=-=-=-=-')

            self.switch_tree_edges(g, h, leaving, entering)
            d = h.span(entering.u, entering.v)
            d = -1 * d + sgn(d)
            self.move_subtree(h, entering.v, d)

            iters += 1
